#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@author:hengk
@contact: hengk@foxmail.com
@datetime:2019-10-25 19:47
"""

import random
import cv2
import string
import barcode
import numpy as np
from io import BytesIO
from barcode.writer import ImageWriter


from libs.utils import load_bgs
from libs.config import load_config
from barcode_renderer.noiser import Noiser
from barcode_renderer.corpus import Corpus

class Renderer(object):
    def __init__(self,flags):
        self.code_min_len_ = flags.min_length
        self.code_max_len_ = flags.max_length
        self.img_w_ = flags.img_width
        self.img_h_ = flags.img_height
        self.bg_dir_ = flags.bg_dir
        self.debug_ = flags.debug
        self.config_ = flags.config_file
        self.corpus_dir_ = flags.corpus_dir
        self.bgs_ = load_bgs(self.bg_dir_)
        self.cfg_ = load_config(self.config_)
        self.noise_ = Noiser(self.cfg_)
        self.corpus_ = Corpus(self.corpus_dir_)

    def gen_bg(self):
        """
        random crop the image
        :param width:
        :param height:
        :return:
        """
        bg = random.choice(self.bgs_)
        assert self.img_w_ > self.img_h_ and bg.shape[1]>self.img_w_ and bg.shape[0]>self.img_h_
        start_x = random.randint(0,bg.shape[1]-self.img_w_-1)
        start_y = random.randint(0,bg.shape[0]-self.img_h_-1)

        out = bg[start_y:start_y+self.img_h_,start_x:start_x+self.img_w_]
        return out

    def gen_code39(self):
        """
        code39编码，数字、大写字母、7个符号
        :param img_index:
        :return:
        """

        number = random.randint(self.code_min_len_,self.code_max_len_)

        chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ-. $/+%"

        idxs = np.random.randint(0, len(chars), size=number)
        code = ''.join([chars[i] for i in idxs])

        ean = barcode.get_barcode_class('code39')
        barcode_instance = ean(code, writer=ImageWriter(),add_checksum=True)

        buffer = BytesIO()
        barcode_instance.write(buffer)

        buffer.seek(0)
        image = np.asarray(bytearray(buffer.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)
        if(self.debug_):
            print("src:"+code)
            print("dst:"+barcode_instance.get_fullcode())
            cv2.imwrite(barcode_instance.get_fullcode()+".jpg",image)

        return image,barcode_instance.get_fullcode()

    def gen_code128(self,code=None):
        #生成条码数字
        number = random.randint(self.code_min_len_, self.code_max_len_)
        chars = string.ascii_letters+string.digits+" !\"#$%&'()*+,-./:;<=>?@[|]~_{}"
        idxs = np.random.randint(0, len(chars), size=number)

        #生成条码图像
        if code is None:
            code = ''.join([chars[i] for i in idxs])
        ean = barcode.get_barcode_class('code128')
        barcode_instance = ean(code, writer=ImageWriter())


        buffer = BytesIO()
        barcode_instance.write(buffer)
        buffer.seek(0)
        image = np.asarray(bytearray(buffer.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        if (self.debug_):
            print("src:" + code)
            print("dst:" + barcode_instance.get_fullcode())
            cv2.imwrite(barcode_instance.get_fullcode() + ".jpg", image)

        return image, barcode_instance.get_fullcode()

    def gen_ean13(self):
        chars = string.digits
        number = 12
        idxs = np.random.randint(0, len(chars), size=number)
        code = ''.join([chars[i] for i in idxs])

        ean = barcode.get_barcode_class('ean13')
        barcode_instance = ean(code, writer=ImageWriter())

        buffer = BytesIO()
        barcode_instance.write(buffer)
        buffer.seek(0)
        image = np.asarray(bytearray(buffer.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        if (self.debug_):
            print("src:" + code)
            print("dst:" + barcode_instance.get_fullcode())
            cv2.imwrite(barcode_instance.get_fullcode() + ".jpg", image)

        return image, barcode_instance.get_fullcode()

    def gen_ean8(self):
        chars = string.digits
        number = 7
        idxs = np.random.randint(0, len(chars), size=number)
        code = ''.join([chars[i] for i in idxs])

        ean = barcode.get_barcode_class('ean8')
        barcode_instance = ean(code, writer=ImageWriter())

        buffer = BytesIO()
        barcode_instance.write(buffer)
        buffer.seek(0)
        image = np.asarray(bytearray(buffer.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        if (self.debug_):
            print("src:" + code)
            print("dst:" + barcode_instance.get_fullcode())
            cv2.imwrite(barcode_instance.get_fullcode() + ".jpg", image)

        return image, barcode_instance.get_fullcode()

    def gen_isbn13(self):
        chars = string.digits
        number = 12
        headers = ['978','97910','97911']
        idxs = np.random.randint(0, len(chars), size=number)
        code = ''.join([chars[i] for i in idxs])
        code = random.choice(headers) + code
        code = code[0:number]

        ean = barcode.get_barcode_class('isbn13')
        barcode_instance = ean(code, writer=ImageWriter())

        buffer = BytesIO()
        barcode_instance.write(buffer)
        buffer.seek(0)
        image = np.asarray(bytearray(buffer.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        if (self.debug_):
            print("src:" + code)
            print("dst:" + barcode_instance.get_fullcode())
            cv2.imwrite(barcode_instance.get_fullcode() + ".jpg", image)

        return image, barcode_instance.get_fullcode()

    def gen_pzn(self):
        chars = string.digits
        number = 6
        checksum = 10
        while checksum==10:
            idxs = np.random.randint(0, len(chars), size=number)
            code = ''.join([chars[i] for i in idxs])
            sum_ = sum([int(x) * int(y) for x, y in enumerate(code, start=2)])
            checksum = sum_ % 11
        ean = barcode.get_barcode_class('pzn')

        barcode_instance = ean(code, writer=ImageWriter())
        buffer = BytesIO()
        barcode_instance.write(buffer)
        buffer.seek(0)
        image = np.asarray(bytearray(buffer.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        if (self.debug_):
            print("src:" + code)
            print("dst:" + barcode_instance.get_fullcode()[-7:])
            cv2.imwrite(barcode_instance.get_fullcode()[-7:] + ".jpg", image)

        return image, barcode_instance.get_fullcode()[-7:]




    def gen_upc(self):
        chars = string.digits
        number = 11
        idxs = np.random.randint(0, len(chars), size=number)
        code = ''.join([chars[i] for i in idxs])

        ean = barcode.get_barcode_class('upc')
        barcode_instance = ean(code, writer=ImageWriter())

        buffer = BytesIO()
        barcode_instance.write(buffer)
        buffer.seek(0)
        image = np.asarray(bytearray(buffer.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        if (self.debug_):
            print("src:" + code)
            print("dst:" + barcode_instance.get_fullcode())
            cv2.imwrite(barcode_instance.get_fullcode() + ".jpg", image)
        return image,barcode_instance.get_fullcode()

    def random_clip(self,image,shrink_ratio=0.03):
        """
        随机裁剪一张图片
        :param image:
        :return:
        """
        h,w,c = image.shape

        clip_w = int(random.uniform(0,shrink_ratio)*w)

        clip_top = int(random.uniform(0,shrink_ratio*3)*h)
        clip_bottom = int(random.uniform(0.25,shrink_ratio*13)*h)

        return image[clip_top:h-clip_bottom,clip_w:w-clip_w]

    def random_perspectivetransform(self,image):
        """
        对图片进行随机透视变换
        :param image:
        :return:
        """
        h,w,_=image.shape

        points = [[0,0],[w-1,0],[0,h-1],[w-1,h-1]]
        for index in range(4):
            offset_x = np.random.uniform(-0.01,0.01)*w
            offset_y = np.random.uniform(-0.01,0.01)*h
            points[index][0] +=offset_x
            points[index][1] +=offset_y
            points[index][0] = max(min(w-1,points[index][0]),0)
            points[index][1] = max(min(h-1,points[index][1]),0)
        points = np.array(points).astype(np.float32)
        points1= np.float32([[0,0],[w-1,0],[0,h-1],[w-1,h-1]])

        M = cv2.getPerspectiveTransform(points, points1)
        image = cv2.warpPerspective(image, M,(w,h))
        return image

    def blend_img(self,background,foreground):
        """
        将前景和背景进行融合
        :param background:
        :param foreground:
        :return:
        """
        fore_h,fore_w,_ = foreground.shape
        back_h,back_w,_ = background.shape
        #使前景比背景大小相同
        background = cv2.resize(background,(fore_w,fore_h),cv2.INTER_AREA)

        #将条码部分画到背景上面
        gray_ground = cv2.cvtColor(foreground,cv2.COLOR_BGR2GRAY)
        _,mask = cv2.threshold(gray_ground,0,255,cv2.THRESH_OTSU|cv2.THRESH_BINARY_INV)
        mask = mask[:,:,np.newaxis].repeat(3,axis=2)
        background[mask==255]=0
        return background


    def gen_img(self, img_index):

        barcode_types = {'code39':self.gen_code39, 'code128':self.gen_code128,
                         'ean13':self.gen_ean13, 'ean8':self.gen_ean8,'pzn':self.gen_pzn,
                         'isbn13':self.gen_isbn13,'upc':self.gen_upc}

        background = self.gen_bg()
        #生成code128类型的条码，主要适用于中国各类快递物流公司
        code = self.corpus_.get_sample(img_index)
        foreground,code =  barcode_types['code128'](code)

        #图片处理
        foreground = self.random_clip(foreground)
        foreground = self.random_perspectivetransform(foreground)
        image = self.blend_img(background, foreground)

        #加入噪声
        image = self.noise_.apply(image)

        return image,code



# a = Renderer(10,20,250,64,"data/bg",False,'configs/default.yaml')
# a.gen_img(1)